#!/usr/bin/env python
# -*- coding: UTF-8 -*-

'''
Copyright (C) 2009 dbzhang800
All rights reserved.
''' 

import sys
import ctypes
import ctypes.wintypes as wintypes
try:
    from PySide import QtCore, QtGui 
except ImportError:
    import sip
    sip.setapi("QString", 2)
    from PyQt4 import  QtCore, QtGui

NULL = 0
ERROR_INSUFFICIENT_BUFFER = 122
STANDARD_RIGHTS_READ = 0x00020000
TOKEN_QUERY = 0x00000008
TOKEN_READ = STANDARD_RIGHTS_READ | TOKEN_QUERY
SECURITY_MAX_SID_SIZE = 68
WinBuiltinAdministratorsSid = 26
TokenElevationType = 18
TokenLinkedToken = 19
TokenElevation = 20
TokenIntegrityLevel = 25
TokenMandatoryPolicy = 27
TOKEN_MANDATORY_POLICY_OFF = 0
SW_SHOWNORMAL = 1
SECURITY_MANDATORY_UNTRUSTED_RID = 0x00000000
SECURITY_MANDATORY_LOW_RID = 0x00001000
SECURITY_MANDATORY_MEDIUM_RID = 0x00002000
SECURITY_MANDATORY_HIGH = 0x00003000
SECURITY_MANDATORY_SYSTEM_RID = 0x00004000
SECURITY_MANDATORY_PROTECTED_PROCESS_RID = 0x00005000

# vista
TokenElevationTypeLimited = 3

kernel32 = ctypes.windll.kernel32
CloseHandle = kernel32.CloseHandle
GetCurrentProcess = kernel32.GetCurrentProcess

shell32 = ctypes.windll.shell32
IsUserAnAdmin = shell32.IsUserAnAdmin
ShellExecuteEx = shell32.ShellExecuteExW

advapi32 = ctypes.windll.advapi32
OpenProcessToken = advapi32.OpenProcessToken
GetTokenInformation = advapi32.GetTokenInformation
CreateWellKnownSid = advapi32.CreateWellKnownSid
CheckTokenMembership = advapi32.CheckTokenMembership
GetSidSubAuthority = advapi32.GetSidSubAuthority
GetSidSubAuthority.restype = ctypes.POINTER(wintypes.DWORD)
GetSidSubAuthorityCount = advapi32.GetSidSubAuthorityCount
GetSidSubAuthorityCount.restype = ctypes.POINTER(ctypes.c_ubyte)

class SID_AND_ATTRIBUTES(ctypes.Structure):
    _fields_ = [("Sid", ctypes.c_void_p),
            ("Attributes", wintypes.DWORD)]

class TOKEN_MANDATORY_LABEL(ctypes.Structure):
    _fields_ = [("Label", SID_AND_ATTRIBUTES)]

class TOKEN_MANDATORY_POLICY(ctypes.Structure):
    _fields_ = [("Policy", wintypes.DWORD)]

def isProcessCanBeElevated():
    canBeElecated = False
    hToken = wintypes.HANDLE()
    OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, ctypes.byref(hToken))
    dwSize = wintypes.DWORD()
    elevationType = wintypes.ULONG() #note, this is enum in c
    GetTokenInformation(hToken, 
            TokenElevationType,
            ctypes.byref(elevationType),
            ctypes.sizeof(elevationType),
            ctypes.byref(dwSize))
    adminSID = ctypes.create_string_buffer(SECURITY_MAX_SID_SIZE)
    dwSize = wintypes.DWORD(SECURITY_MAX_SID_SIZE)
    CreateWellKnownSid(WinBuiltinAdministratorsSid,
            NULL,
            ctypes.byref(adminSID),
            ctypes.byref(dwSize))
    isAdmin = wintypes.BOOL(0)
    if elevationType.value == TokenElevationTypeLimited:
        hUnfilteredToken = wintypes.HANDLE()
        GetTokenInformation(hToken, 
                TokenLinkedToken,
                ctypes.byref(hUnfilteredToken),
                ctypes.sizeof(wintypes.HANDLE),
                ctypes.byref(dwSize))
        if CheckTokenMembership(hUnfilteredToken, ctypes.byref(adminSID), ctypes.byref(isAdmin)):
            if isAdmin.value:
                canBeElecated = True
        CloseHandle(hUnfilteredToken)
    else:
        isAdmin.value = IsUserAnAdmin()
    CloseHandle(hToken)
    return isAdmin.value!=0, canBeElecated

def isProcessElevated():
    class TOKEN_ELEVATION(ctypes.Structure):
        _fields_ = [("TokenIsElevated", wintypes.DWORD)]

    isElevated = False
    hToken = wintypes.HANDLE()
    OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, ctypes.byref(hToken))
    dwSize = wintypes.DWORD()
    elevation = TOKEN_ELEVATION()
    if not GetTokenInformation(hToken, TokenElevation, ctypes.byref(elevation),
            ctypes.sizeof(elevation), ctypes.byref(dwSize)):
        return "NA"
    return elevation.TokenIsElevated != 0


def getProcessIntegrityLevel():
    hToken = wintypes.HANDLE()
    OpenProcessToken(GetCurrentProcess(), TOKEN_READ, ctypes.byref(hToken))
    isReturn = False

    # get integrity level
    integrityLevel = wintypes.DWORD()
    dwNeededSize = wintypes.DWORD()
    if not GetTokenInformation(hToken, TokenIntegrityLevel, NULL, 0, ctypes.byref(dwNeededSize)):
        tokenInfo = TOKEN_MANDATORY_LABEL()
        if ctypes.GetLastError() == ERROR_INSUFFICIENT_BUFFER:
            ctypes.resize(tokenInfo, dwNeededSize.value)
            if GetTokenInformation(hToken, TokenIntegrityLevel, ctypes.byref(tokenInfo),
                    dwNeededSize, ctypes.byref(dwNeededSize)):
                nSize = GetSidSubAuthorityCount(tokenInfo.Label.Sid).contents.value - 1
                integrityLevel = GetSidSubAuthority(tokenInfo.Label.Sid, nSize).contents
                #print integrityLevel.value
                isReturn = True

    # get policy if the integrity level available
    policy = TOKEN_MANDATORY_POLICY(TOKEN_MANDATORY_POLICY_OFF)
    if isReturn:
        dwNeededSize = wintypes.DWORD(ctypes.sizeof(policy))
        GetTokenInformation(hToken, TokenMandatoryPolicy, ctypes.byref(policy),
                dwNeededSize, ctypes.byref(dwNeededSize))
    CloseHandle(hToken)

    stringDict = {SECURITY_MANDATORY_UNTRUSTED_RID: "Untrusted",
            SECURITY_MANDATORY_LOW_RID: "Low",
            SECURITY_MANDATORY_MEDIUM_RID: "Medium",
            SECURITY_MANDATORY_HIGH: "High",
            SECURITY_MANDATORY_SYSTEM_RID: "System",
            SECURITY_MANDATORY_PROTECTED_PROCESS_RID: "Protected"}

    return stringDict[integrityLevel.value] if isReturn else "NA"

def restartElevateProcess():
    class SHELLEXECUTEINFO(ctypes.Structure):
        class U(ctypes.Union):
            _fields_ = [("hIcon", wintypes.HANDLE),
                    ("hMonitor", wintypes.HANDLE)]

        _fields_ = [("cbSize", wintypes.DWORD),
                ("fMask", wintypes.ULONG),
                ("hwnd", wintypes.HWND),
                ("lpVerb", ctypes.POINTER(ctypes.c_wchar)),
                ("lpFile", ctypes.POINTER(ctypes.c_wchar)),
                ("lpParameters", ctypes.POINTER(ctypes.c_wchar)),
                ("lpDirectory", ctypes.POINTER(ctypes.c_wchar)),
                ("nShow", ctypes.c_int),
                ("hInstApp", wintypes.HINSTANCE),
                ("lpIDList", ctypes.c_void_p),
                ("lpClass", ctypes.POINTER(ctypes.c_wchar)),
                ("hkeyClass", wintypes.HKEY),
                ("dwHotKey", wintypes.DWORD),
                ("DUMMYUNIONNAME", U),
                ("hProcess", wintypes.HANDLE)]

    sei = SHELLEXECUTEINFO()
    sei.cbSize = ctypes.sizeof(SHELLEXECUTEINFO)
    sei.lpVerb = ctypes.create_unicode_buffer("runas")
    sei.lpFile = ctypes.create_unicode_buffer(QtGui.qApp.applicationFilePath())
    sei.lpParameters = ctypes.create_unicode_buffer(" ".join(sys.argv))
    sei.nShow = SW_SHOWNORMAL

    if ShellExecuteEx(ctypes.byref(sei)):
        QtGui.qApp.quit() # quit itself

class Window(QtGui.QWidget):
    def __init__(self,  parent=None):
        super(Window,  self).__init__(parent)
        self.setWindowTitle("SelfElevation")
        self.elevateButton = QtGui.QPushButton("Self elevate", self)
        self.elevateButton.setIcon(QtGui.qApp.style().standardIcon(QtGui.QStyle.SP_VistaShield))
        isAdmin, canElevated = isProcessCanBeElevated()
        font = QtGui.qApp.font()
        font.setBold(True)
        font.setPointSize(font.pointSize()+2)
        QtGui.qApp.setFont(font)
        
        vbox = QtGui.QVBoxLayout(self)
        vbox.addWidget(QtGui.QLabel("{:20s}{:s}".format("IsUserInAdminGroup:", str(isAdmin))))
        vbox.addWidget(QtGui.QLabel("{:20s}{:s}".format("IsProcessElevated:", str(isProcessElevated()))))
        vbox.addWidget(QtGui.QLabel("{:20s}{:s}".format("IntegrityLevel:", getProcessIntegrityLevel())))
        vbox.addWidget(self.elevateButton)
        self.setLayout(vbox)
        self.elevateButton.setVisible(canElevated)
        self.elevateButton.clicked.connect(self.onElevateButtonClicked)

    def onElevateButtonClicked(self):
        restartElevateProcess()


if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    w = Window()
    w.show()
    sys.exit(app.exec_())
